from openai import OpenAI
import streamlit as st
import json
import os
import glob
import re
import shutil  

#欢迎访问ManyiAPI聚合接口站：https://api.manyi88.top，注册优惠

# 示例：模型和相应的 API 密钥
default_key = "***"  # 默认令牌
sale_key = "***" #自定义默认1.0
guan_key ="***" #管转令牌3倍
az_key="***" #纯AZ，1.5倍
claude_key="***" # claude 8倍
guan5_key ="***" #管转令牌5倍
model_keys = {
    "gpt-4o-mini-2024-07-18": az_key,  #1.5
    "o1-mini": guan5_key,    
    "gpt-4o-2024-08-06": az_key,
    "claude-3-5-sonnet-20240620": claude_key
}

st.title("ChatGPT-like Clone")

selected_model = st.sidebar.selectbox("选择模型", list(model_keys.keys()))
api_key = model_keys[selected_model]
api_url = "https://api.manyi88.top/v1"   #欢迎访问ManyiAPI聚合接口站：https://api.manyi88.top，注册优惠

client = OpenAI(api_key=api_key, base_url=api_url)

data_dir = "data"
backup_dir = "data_bak"
if not os.path.exists(data_dir):
    os.makedirs(data_dir)
if not os.path.exists(backup_dir):
    os.makedirs(backup_dir)

session_id_file = os.path.join(data_dir, "session_id.txt")

def load_session_id():
    if os.path.exists(session_id_file):
        with open(session_id_file, "r") as f:
            return int(f.read().strip())
    return 0

def save_session_id(session_id):
    with open(session_id_file, "w") as f:
        f.write(str(session_id))

if "openai_model" not in st.session_state:
    st.session_state["openai_model"] = selected_model

def load_history(file_path):
    match = re.search(r'chat_history_(\d+)\.json', os.path.basename(file_path))
    if match:
        st.session_state.session_id = int(match.group(1))

    with open(file_path, "r") as f:
        st.session_state.messages = json.load(f)

def load_latest_history():
    history_files = sorted(glob.glob(os.path.join(data_dir, "*.json")), key=os.path.getmtime)
    if history_files:
        latest_file = history_files[-1]
        load_history(latest_file)

if "messages" not in st.session_state:
    st.session_state.messages = []
    load_latest_history()

if "session_id" not in st.session_state:
    st.session_state.session_id = load_session_id()
    

def save_history():
    file_path = os.path.join(data_dir, f"chat_history_{st.session_state.session_id}.json")
    try:
        json.dumps(st.session_state.messages)
        with open(file_path, "w", encoding="utf-8") as f:
            json.dump(st.session_state.messages, f, ensure_ascii=False)
    except (TypeError, ValueError):
        st.error("数据格式不正确，无法保存记录")

def move_history(file_path):
    # Move the history file to the backup directory
    if os.path.exists(file_path):
        destination_file = os.path.join(backup_dir, os.path.basename(file_path))
        shutil.move(file_path, destination_file)

def delete_history(file_path):
    if os.path.exists(file_path):
        os.remove(file_path)

def get_chat_title(messages):
    if messages:
        first_message = messages[0]["content"]
        title = first_message if first_message else "空的聊天"
        return title[:12]
    return "空的聊天"[:12]

def show_history_files(page=0, page_size=10):
    history_files = [(f, os.path.getmtime(f)) for f in glob.glob(os.path.join(data_dir, "*.json"))]
    history_files.sort(key=lambda x: x[1], reverse=True)

    total_files = len(history_files)
    total_pages = (total_files // page_size) + (1 if total_files % page_size > 0 else 0)

    start = page * page_size
    end = start + page_size
    display_files = history_files[start:end]

    with st.sidebar.expander("历史聊天记录"):
        for index, (file_path, _) in enumerate(display_files):
            with open(file_path, "r") as f:
                chat_history = json.load(f)

            file_name = get_chat_title(chat_history)

            col1, col2, col3 = st.sidebar.columns([4, 1, 1])  # Create three columns
            with col1:
                if st.button(file_name, key=f"load_{index}"):
                    load_history(file_path)
                    st.rerun()

            with col2:
                if st.button("📦", key=f"move_{index}",help=f"移动到备份文件夹"):
                    move_history(file_path)
                    st.success(f"{file_name} 已移动到备份。")
                    st.rerun()

            with col3:
                if st.button("❌", key=f"delete_{index}",help=f"删除"):
                    delete_history(file_path)
                    st.success(f"{file_name} 已删除。")
                    st.rerun()

        # Pagination
        if page > 0:
            if st.button("上一页"):
                st.session_state.current_page -= 1
                st.rerun()

        if page < total_pages - 1:
            if st.button("下一页"):
                st.session_state.current_page += 1
                st.rerun()

if "current_page" not in st.session_state:
    st.session_state.current_page = 0

for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

output_mode = st.sidebar.selectbox("选择输出模式", ["流式输出 (Stream)", "非流式输出 (Non-stream)"])

# 历史数量输入框
# 初始化 session_state 中的历史数量
if "history_count" not in st.session_state:
    st.session_state.history_count = 0

# 在侧边栏添加滑动条
if len(st.session_state.messages) == 0:
    max_history_count=10
else:
    max_history_count = len(st.session_state.messages)

st.sidebar.slider(
    "选择使用的历史消息数量(共" + str(len(st.session_state.messages)) + "条)",
    min_value=0,
    max_value=max_history_count,
    value=st.session_state.history_count,  # 默认值
    key="history_count"  # 使用一个唯一的键来存储选择值
)

# 显示当前选择的历史消息数量
st.sidebar.write(f"您选择的历史消息数量是: {st.session_state.history_count}")


prompt = st.chat_input("What is up?")
if prompt:
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    with st.chat_message("assistant"):
        client.api_key = model_keys[selected_model]

        stream = output_mode == "流式输出 (Stream)"

        try:
            if st.session_state.history_count > 0:
                messages_to_send = (
                    [{"role": m["role"], "content": m["content"]}
                     for m in st.session_state.messages[-st.session_state.history_count:]]
                )
            else:
                messages_to_send = [{"role": "user", "content": prompt}]

            res = client.chat.completions.create(
                model=selected_model,
                messages=messages_to_send,
                stream=stream,
            )

            if stream:
                assistant_message = st.write_stream(res)
                if assistant_message:
                    st.session_state.messages.append({"role": "assistant", "content": assistant_message})
                else:
                    st.warning("收到空响应。")
            else:       
                if len(res.choices) > 0:
                    assistant_message = res.choices[0].message.content
                    if assistant_message:
                        st.markdown(assistant_message)
                        st.session_state.messages.append({"role": "assistant", "content": assistant_message})
                    else:
                        st.warning("收到空响应。")
                else:
                    st.warning("响应格式不正确，未找到有效消息。")

            save_history()
        except Exception as e:
            st.error(f"发生错误：{e}")

st.sidebar.header("操作")
if st.sidebar.button("New Chat"):
    st.session_state.messages = []  # Clear current session
    st.session_state.session_id = load_session_id()
    st.session_state.session_id += 1
    save_session_id(st.session_state.session_id)
    st.success("当前会话已清空。")

st.sidebar.header("历史对话")
show_history_files(st.session_state.current_page)

